import tensorflow as tf
from tensorflow.keras import layers,losses,optimizers
from tensorflow.keras.datasets import mnist
import matplotlib.pyplot as plt
import numpy as np
# 使用sequential 搭建两层循环神经网络
# 自动管理记忆层

(x_train, y_train), (x_test, y_test)=mnist.load_data()

model = tf.keras.Sequential([
    # return_sequences=True state内部记忆参数传递至下一层
    layers.LSTM(64, dropout=0.3, return_sequences=True),
    layers.LSTM(10),
    layers.Softmax(),
])

model.compile(
    optimizer=optimizers.Adam(),
    loss=losses.SparseCategoricalCrossentropy(),
    metrics=['accuracy']
)

model.fit(tf.cast(x_train,dtype=tf.float32),tf.cast(y_train,dtype=tf.float32),epochs=10)

r = np.random.randint(0, 10000)

x_test_img = tf.reshape(x_test[r],[28,28])
plt.imshow(x_test_img)
plt.show()

pred = model.predict(tf.cast(x_test[r:r+1], dtype=tf.float32))
print('pred_y=',np.argmax(pred,1))
print('ture_y=',y_test[r:r+1])




